import time

def loadDataSet():
    return [[1, 3, 4], [2, 3, 5], [1, 2, 3, 5], [2, 5]]


def loadDataSet1():
    fr = open('mushroom.dat')
    dataSet = []
    for line in fr.readlines():
        currLine = []
        lineArr = line.strip().split()
        for num in lineArr:
            currLine.append(int(num))
        dataSet.append(currLine)
    return dataSet


# def loadDataSet2():
#     dataMat = []
#     f = open("mushroom.dat")
#     for line in f.readlines():
#         # strip()的用法是移除字符串头尾指定的字符序列。
#         # 返回移除字符串头尾指定的字符生成的新字符串
#         lineArr = line.strip().split()
#         dataMat.append(lineArr)
#     return dataMat


def createC1(dataSet):
    """
    初始化：构造候选1项集
    :return:
    """
    C1 = []
    for transaction in dataSet: # 对于数据集中的每一个记录
        for item in transaction: # 对于记录中的每一项
            if not [item] in C1:
                C1.append([item])
    C1.sort()
    return list(map(frozenset, C1)) # map() 会根据提供的函数对指定序列做映射。第一个参数 function 以参数序列中的每一个元素调用 function 函数，返回包含每次 function 函数返回值的新列表。
                                   # Python 3.x 返回迭代器。
                                   # 一定要加list!!!!


def scanD(D, Ck, minSupport):
    """
    通过候选k项集Ck来构造频繁k项集Lk
    :param D: 数据集
    :param Ck: 候选k项集
    :param minSupport: 阈值
    :return:
    """
    ssCnt = {} # 用来记录每一个候选项集出现的次数的字典
    # 两个for循环用来计算候选k项集
    for tid in D: # 对于数据集D中的每一个记录
        for can in Ck: # 对于候选项集中的每一个候选项
            if can.issubset(tid): # issubset() 方法用于判断集合的所有元素是否都包含在指定集合中，如果是则返回 True，否则返回False
                if can not in ssCnt:
                    ssCnt[can] = 1
                else:
                    ssCnt[can] += 1
    numItems = float(len(D)) # 计算数据集的大小，用于计算支持度
    retList = [] # 用于存储频繁k项集
    supportData = {} # 用来存储各个候选项集的支持度
    # for循环寻找频繁k项集
    for key in ssCnt:
        support = ssCnt[key] / numItems # 支持度的计算
        if support >= minSupport:
            retList.insert(0, key) # insert() 函数用于将指定对象插入列表的指定位置。
        supportData[key] = support
    return retList, supportData


def aprioriGen(Lk, k):
    """

    :param Lk: 频繁k项集
    :param k: 项集元素个数，即k项集
    :return:
    """
    retList = [] # 存储k项集通过连接形成的候选k+1项集
    lenLk = len(Lk)
    for i in range(lenLk):
        for j in range(i+1, lenLk):
            L1 = list(Lk[i])[:k-2]
            L2 = list(Lk[j])[:k-2]
            L1.sort()
            L2.sort()
            if L1 == L2:
                retList.append(Lk[i] | Lk[j])
    return retList


def apriori(dataSet, minSupport=0.5):
    C1 = createC1(dataSet)
    D = list(map(set, dataSet)) # 集合中的元素不重复
    L1, supportData = scanD(D, C1, minSupport)
    L = [L1]
    k = 2
    while (len(L[k-2]) > 0):
        Ck = aprioriGen(L[k-2], k)
        Lk, supK = scanD(D, Ck, minSupport)
        supportData.update(supK)
        L.append(Lk)
        k += 1
    return L, supportData


if __name__ == '__main__':
    startTime = time.time()
    dataSet = loadDataSet1()
    # print(dataSet)
    L, supportData = apriori(dataSet, 0.5)
    endTime = time.time()
    k = 1
    num = 0
    for frequentk in L:
        if len(frequentk) == 0:
            break
        print("频繁%d项集：" % k, end="\t")
        for item in frequentk:
            print(item, end="\t")
            num += 1
        print()
        k += 1
    print("共发现频繁情节：%d个" % num)
    print("程序耗时：%f" % (endTime - startTime))